package geppetto.phraseHMM;


import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentSymetrization;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.DynamicAlignmentSet;
import geppetto.cat.alignments.AlignmentEvaluator.Evaluation;
import geppetto.cat.alignments.output.AlignerOutputLatex;
import geppetto.cat.alignments.phrasal.PhrasalHmmAgreementModel;
import geppetto.cat.alignments.phrasal.PhrasalHmmModel;
import geppetto.cat.alignments.phrasal.PhraseDistortionTable;
import geppetto.cat.alignments.phrasal.PhraseTranslationTable;
import geppetto.cat.common.StaticTools;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.SparseTranslationTable;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintStream;



/** class to train and save an (HMM) alignment model */
public class PhraseAlignmentsExtraction {

	public static void main(String[] args) throws IOException {
		String corpusFile = args[0];
		int size = Integer.parseInt(args[1]); // 100k
		int maxSentenceSize = Integer.parseInt(args[2]); // 40
		int numberIterations = Integer.parseInt(args[3]); // 5
		boolean trainWithResults = Boolean.parseBoolean(args[4]);
		int numberIterationsWithResults = Integer.parseInt(args[5]);
		double smooth = Double.parseDouble(args[6]);
		String dir = args[7];
		String modelName = args[8];
		String wordAlignmentFile = args[9];
		String wordAlignmentFileType = args[10];
		int sourceMaxSize = Integer.parseInt(args[11]);
		int targetMaxSize = Integer.parseInt(args[12]);
		String initialPhraseTableFile = args[13];
		String initialPhraseTableFileRev = args[14];
		String initialNullPhraseFile = args[15];
		String initialNullPhraseFileRev = args[16];
		String initialDistortionTableDirectory = args[17];
		String initialDistortionTableDirectoryRev = args[18];
		
		System.out.println("Saving Models experiment: ");
		System.out.println("Corpus " + corpusFile);
		System.out.println("Size " + size);
		System.out.println("Smoothing " + smooth);
		System.out.println("Max Sentence size " + maxSentenceSize);
		System.out.println("Number of EM iterations " + numberIterations);
		System.out.println("Train with results " + trainWithResults + " iter "+ numberIterationsWithResults);
		System.out.println("OutputDir " + dir);
		System.out.println("Model Type " + modelName);
		System.out.println("Word alignments file " + wordAlignmentFile);
		System.out.println("Word alignments file type " + wordAlignmentFileType);
		System.out.println("phrase source max size " + sourceMaxSize);
		System.out.println("phrase target max size " + targetMaxSize);
		System.out.println("phrase table file " + initialPhraseTableFile);
		System.out.println("phrase reverse table file " + initialPhraseTableFileRev);
		System.out.println("phrase distortion dir " + initialDistortionTableDirectory);
		System.out.println("phrase distortion dir " + initialDistortionTableDirectoryRev);
		System.out.println("---");
		BilingualCorpus corpusF = BilingualCorpus.getCorpusFromFileDescription(
				corpusFile, size, maxSentenceSize);
		BilingualCorpus corpusB = corpusF.reverse();
		String baseDir = dir + "/" + modelName + "/" + corpusF.getName() + "/" + size;
		String modelDir = baseDir + "/model/";
		StaticTools.createDir(modelDir);
		System.out.println("Saving models to "+ modelDir);
		
		//loading word alignments for base
		AlignmentsSet alignments = buildAlignmentSet(wordAlignmentFileType, wordAlignmentFile, corpusF);
		
		//loading wordtries
		WordTrie sourceVocab = loadSourceWordTrie(sourceMaxSize, corpusF);
		WordTrie targetVocab = loadTargetWordTrie(targetMaxSize, corpusF);
		
		PhraseTranslationTable translationTable = new PhraseTranslationTable(corpusF, sourceVocab, targetVocab, sourceMaxSize, targetMaxSize);
		if(!initialPhraseTableFile.equals("-")){
			translationTable.loadFromPhraseTableFile(initialPhraseTableFile);
		}
		if(!initialNullPhraseFile.equals("-")){
			translationTable.loadFromNullPhraseTable(initialNullPhraseFile);
		}		
		PhraseDistortionTable distortionTable = new PhraseDistortionTable();
		if(!initialDistortionTableDirectory.equals("-")){
			if(new File(initialDistortionTableDirectory).isDirectory()){
				distortionTable.loadDistortionTableFromDir(initialDistortionTableDirectory);				
			}
			else{
				distortionTable.loadDistortionTableFromFile(initialDistortionTableDirectory);
			}
		}
		PhrasalHmmModel model = new PhrasalHmmModel(corpusF, sourceVocab, targetVocab, alignments, translationTable, distortionTable, sourceMaxSize, targetMaxSize);

		PhraseTranslationTable translationTableRev = new PhraseTranslationTable(corpusB, targetVocab, sourceVocab, targetMaxSize, sourceMaxSize);
		if(!initialPhraseTableFileRev.equals("-")){
			translationTableRev.loadFromPhraseTableFile(initialPhraseTableFileRev);
		}
		if(!initialNullPhraseFileRev.equals("-")){
			translationTableRev.loadFromNullPhraseTable(initialNullPhraseFileRev);
		}
		PhraseDistortionTable distortionTableRev = new PhraseDistortionTable();
		if(!initialDistortionTableDirectoryRev.equals("-")){
			if(new File(initialDistortionTableDirectoryRev).isDirectory()){
				distortionTableRev.loadDistortionTableFromDir(initialDistortionTableDirectoryRev);
			}
			else{
				distortionTableRev.loadDistortionTableFromFile(initialDistortionTableDirectoryRev);
			}
		}
		PhrasalHmmModel revModel = new PhrasalHmmModel(corpusB, targetVocab, sourceVocab, alignments.reverse(), translationTableRev, distortionTableRev, targetMaxSize, sourceMaxSize);
		
		PhrasalHmmAgreementModel agreementModel = new PhrasalHmmAgreementModel(model,revModel);
		
		PrintStream phraseExtractionFile = new PrintStream(new File(modelDir+"/phrases.txt"), "UTF8");
		PrintStream nullPhraseExtractionFile = new PrintStream(new File(modelDir+"/nullphrases.txt"), "UTF8");
		PrintStream reorderingExtractionFile = new PrintStream(new File(modelDir+"/reordering.txt"), "UTF8");
		PrintStream phraseExtractionFileB = new PrintStream(new File(modelDir+"/phrasesB.txt"), "UTF8");
		PrintStream nullPhraseExtractionFileB = new PrintStream(new File(modelDir+"/nullphrasesB.txt"), "UTF8");
		PrintStream reorderingExtractionFileB = new PrintStream(new File(modelDir+"/reorderingB.txt"), "UTF8");
		agreementModel.extractPhrases(phraseExtractionFile, nullPhraseExtractionFile, reorderingExtractionFile, phraseExtractionFileB, nullPhraseExtractionFileB, reorderingExtractionFileB);		
		
		System.out.println("Done extracting");
	}

	public static AlignmentsSet buildAlignmentSet(String alignmentsFileType, String alignmentFile, BilingualCorpus corpus) throws FileNotFoundException, IOException {
		if(alignmentsFileType.equals("posterior")){
			return DynamicAlignmentSet.getAlignmentFromPosterior(alignmentFile, BilingualCorpus.TRAIN_CORPUS, corpus);
		}
		else if (alignmentsFileType.equals("moses")){
			return DynamicAlignmentSet.getAlignmentFromMoses(alignmentFile, BilingualCorpus.TRAIN_CORPUS, corpus);
		}
		return null;
	}
	
	public static WordTrie loadSourceWordTrie(int maxLength, BilingualCorpus corpus) throws FileNotFoundException, IOException{
		WordTrie sourceVocab = new WordTrie(maxLength); 		
		sourceVocab.addPhrasesAllSentences(corpus._trainSourceSentences);
		sourceVocab.compactTrie();
		return sourceVocab;
	}

	public static WordTrie loadTargetWordTrie(int maxLength, BilingualCorpus corpus) throws FileNotFoundException, IOException{
		WordTrie targetVocab = new WordTrie(maxLength); 
		targetVocab.addPhrasesAllSentences(corpus._trainForeignSentences);
		targetVocab.compactTrie();
		return targetVocab;
	}

}
